import platform
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import argostranslate.package
import argostranslate.translate
from fastapi import FastAPI
from fastapi import HTTPException, Depends
from fastapi import Header
from pydantic import BaseModel
from starlette import status

from app.log import x_logger

# 你定义一个固定 token，比如
API_TOKEN = "996.icu"


def verify_token(authorization: str = Header(...)):
    if not authorization.startswith("Bearer "):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authorization header"
        )
    token = authorization[7:]  # 去掉 "Bearer "
    if token != API_TOKEN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or missing token"
        )


auth_required = Depends(verify_token)


# 初始化和加载模型
def install_argos_models(models_dir):
    models_dir = Path(models_dir)
    installed = argostranslate.translate.get_installed_languages()
    installed = [l.code for l in installed]
    for model_path in models_dir.glob("*.argosmodel"):
        # 解析语言对（如 zh_en-1_9.argosmodel）
        name = model_path.stem  # 文件名不含后缀
        parts = name.split('-')
        if len(parts) >= 3:
            lang_pair = parts[1]  # 比如 en_zh
            from_lang, to_lang = lang_pair.split('_')
            if all(l in installed for l in [from_lang, to_lang]):
                x_logger.info(f"Skip: {from_lang}->{to_lang} already installed")
                continue
        x_logger.info(f"Installing: {model_path}")
        argostranslate.package.install_from_path(Path(model_path))  # 用str，确保兼容


# 启动时加载模型
if platform.system() == "Windows":
    models_dir = r"D:\WorkSpace\dev\argostranslate-models"
else:
    models_dir = r"/home/argos/models"

install_argos_models(models_dir)
installed_languages = argostranslate.translate.get_installed_languages()


def get_translation_inst(from_code='zh', to_code='en'):
    from_lang = next((l for l in installed_languages if l.code == from_code), None)
    to_lang = next((l for l in installed_languages if l.code == to_code), None)
    if from_lang and to_lang:
        return from_lang.get_translation(to_lang)
    return None


app = FastAPI(
    docs_url="/docs",  # Swagger UI 路径，默认 "/docs"
    redoc_url="/redoc",  # Redoc 路径，默认 "/redoc"
    openapi_url="/openapi.json"  # OpenAPI schema 路径
)


class TranslateRequest(BaseModel):
    contents: list
    from_lang: str = "zh"
    to_lang: str = "en"


class TranslateResponse(BaseModel):
    results: list


@app.post("/translate", response_model=TranslateResponse)
def translate(req: TranslateRequest, _: None = auth_required):
    start_time = time.time()
    x_logger.info(f"Received translation items: {len(req.contents)}")
    translation = get_translation_inst(req.from_lang, req.to_lang)
    if not translation:
        raise HTTPException(status_code=400, detail="Model not found for language pair")
    with ThreadPoolExecutor() as executor:
        results = list(executor.map(translation.translate, req.contents))
    x_logger.info(f"Translated {len(req.contents)} items in {time.time() - start_time:.2f} seconds")
    return {"results": results}
